package com.xbai.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, LongType, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

/**
  * 用户自定义聚合函数
  * 强类型的 Dataset 和弱类型的 DataFrame 都提供了相关的聚合函数，
  * 如 count()，countDistinct()，avg()，max()，min()。除此之外，用户可以设定自己的自定义聚合函数。
 *
  * 弱类型用户自定义聚合函数：通过继承 UserDefinedAggregateFunction 来实现用户自定义聚合函数。
  * 案例：求平均年龄的自定义聚合函数。
  *
  * @author xbai
  * @Date 2021/1/5
  */
object SparkSQL05_UDAF {

  def main(args: Array[String]): Unit = {
    // 创建配置对象
    val conf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("udaf")
    // 创建 SparkSQL 的环境对象
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()

    // 在进行转换之前，需要引入隐式转换规则

    val udaf = new MyAgeAvgFunction

    spark.udf.register("avgAge", udaf)

    // 读取数据，构建 DataFrame
    val df: DataFrame = spark.read.json("in/user.json")

    // 将 DataFrame 转换位一张表
    df.createOrReplaceTempView("people")

    spark.sql("select avgAge(age) from people").show()

    spark.stop()
  }
}

/**
  * 声明用户自定义聚合函数
  * 1.继承 UserDefinedAggregateFunction
  * 2.实现方法
  */
class MyAgeAvgFunction extends UserDefinedAggregateFunction{

  // 函数输入的数据结构
  override def inputSchema: StructType = {
    new StructType().add("age", LongType)
  }

  // 计算时的数据结果
  override def bufferSchema: StructType = {
    new StructType().add("sum", LongType).add("count", LongType)
  }

  // 函数返回的数据类型
  override def dataType: DataType = DoubleType

  // 函数是否稳定
  override def deterministic: Boolean = true

  // 计算之前的缓冲区的初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0L
    buffer(1) = 0L
  }

  // 根据查询结果更新缓冲区数据
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    buffer(0) = buffer.getLong(0) + input.getLong(0)
    buffer(1) = buffer.getLong(1) + 1
  }

  // 将多个节点的缓冲区合并
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    // sum
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
    // count
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  // 计算
  override def evaluate(buffer: Row): Any = buffer.getLong(0).toDouble / buffer.getLong(1)
}